/*
 * sapphire-backend
 *
 * Copyright (C) 2018 Alyssa Rosenzweig
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02111-1301  USA
 *
 */

#include <stdio.h>
#include <stdint.h>
#include <assert.h>

#include <glib.h>
#include <libsoup/soup.h>

#include <signal.h>
#include <string.h>
#ifndef _WIN32
#include <unistd.h>
#else
#include "win32/win32dep.h"
#endif
#include <fcntl.h>
#include <gio/gio.h>
#include <gio/gunixsocketaddress.h>

#include "json_compat.h"

/* Represents a connection to a given client. Minimal state should be kept
 * here, since connections are device-specific, not for the client as a whole.
 * Essentially, just enough for the websocket metadata and a little potpourrie */

typedef struct Connection {
	/* The reference counted connection itself */
	SoupWebsocketConnection *connection;

	/* The reference counted proxied connection */
	GSocketConnection *proxy_connection;
	GDataInputStream *distream;

	/* Has this connection authenticated yet? */
	gboolean is_authenticated;

	/* Remote IP Address. Must be g_free'd on destroy */
	gchar *ip_address;
} Connection;

static void
sapphire_close_connection(Connection *conn);

/* Whether this proxy supports multi-user mode */

static gboolean multi_user = FALSE;

/* Helper to serialize and broadcast */

#define WS_PORT 7070

/* List of all authenticated connections. These connections will be broadcasted
 * to by broadcast_raw_packet. */

GList *authenticated_connections = NULL;

static void
sapphire_send_raw_packet(Connection *conn, const char *packet)
{
	if (soup_websocket_connection_get_state(conn->connection) != SOUP_WEBSOCKET_STATE_OPEN) {
		printf("Ignoring write to closed websocket\n");
		return;
	}

	soup_websocket_connection_send_text(conn->connection, packet);
}

static void
sapphire_send(Connection *conn, JsonObject *msg)
{
	gchar *str = json_object_to_string(msg);
	sapphire_send_raw_packet(conn, str);
	g_free(str);
}

static JsonNode *
json_parse_to_root(gchar *frame)
{
	JsonParser *parser = json_parser_new();

	if (!json_parser_load_from_data(parser, frame, -1, NULL)) {
		fprintf(stderr, "Error parsing response: %s\n", frame);
		return NULL;
	}

	return json_parser_get_root(parser);
}

/* Serialize accounts to/from disk. TODO: Serializing back */

typedef struct {
	gchar *password_hash;
} SapphireAccount;

GHashTable *username_to_account;

static void
sapphire_load_accounts(const gchar *name)
{
	GError *error = NULL;
	gchar *contents = NULL;
	gboolean success = g_file_get_contents(name, &contents, NULL, &error); 

	if (!success) {
		printf("Bad read\n");
		exit(1);
	}

	JsonNode *root = json_parse_to_root(contents);
	g_free(contents);

	if (root == NULL) {
		printf("NULL root, ignoring\n");
		return;
	}

	JsonArray *accounts = json_node_get_array(root); 
	int len = json_array_get_length(accounts);

	for (int i = 0; i < len; ++i) {
		JsonObject *account = json_array_get_object_element(accounts, i);
		const gchar *name = json_object_get_string_member(account, "name");
		SapphireAccount *sapph = g_new0(SapphireAccount, 1);
		sapph->password_hash = g_strdup(json_object_get_string_member(account, "passwordHash"));

		g_hash_table_insert(username_to_account, g_strdup(name), sapph);
	}

	if (len > 1 && !multi_user) {
		printf("ERROR: Multi-user mode not enabled (did you enable TLS?) but multiple accounts in the database\n");
		exit(1);
	}
}



/* Functions for icon proxying, TODO segregate by user */

GHashTable *username_to_icon;

/* Generic container for icons to paper over the difference between
 * PurpleStoredImage and PurpleBuddyIcon */

typedef struct {
	const gchar *extension;
	size_t size;
	gconstpointer data;
} SapphireIcon;

static void
soup_icon_callback(SoupServer *server,
                       SoupMessage *msg,
                       const char *path,
                       GHashTable *query,
                       SoupClientContext *client,
                       gpointer user_data)
{
	if (msg->method != SOUP_METHOD_GET) {
		soup_message_set_status (msg, SOUP_STATUS_NOT_IMPLEMENTED);
		return;
	}

	/* Extract the name from the query */

	if (!g_hash_table_contains(query, "name")) {
		soup_message_set_status (msg, SOUP_STATUS_NOT_FOUND);
		return;
	}

	const gchar *name = g_hash_table_lookup(query, "name");

	/* Search for icon */
	SapphireIcon *icon = g_hash_table_lookup(username_to_icon, name);

	if (!icon) {
		soup_message_set_status (msg, SOUP_STATUS_NOT_FOUND);
		return;
	}

	/* We found it, so return the icon appropriately */

	gchar *mimetype = g_strdup_printf("image/%s", icon->extension);

	if (!icon->data) {
		soup_message_set_status (msg, SOUP_STATUS_NOT_FOUND);
		return;
	}

	soup_message_set_status(msg, SOUP_STATUS_OK);
	soup_message_set_response(msg, mimetype, SOUP_MEMORY_TEMPORARY, icon->data, icon->size);

	/* Set cacheing header */
	soup_message_headers_append(msg->response_headers, "Cache-Control", "public, max-age=2592000");

	g_free(mimetype);
}

/* IP -> rate limit hash table. Rate limits are integers of the number of
 * seconds rate limited. If positive, the rate limit is active. If negative, it
 * is inactive and merely stored for posterity. Rate limits are mostly reset on
 * a successful authentication */

GHashTable *rate_limits;

/* Callback for a timer once the rate limit is restored for the next attempt */

static gboolean
sapphire_restore_rate_limit(gpointer user_data)
{
	gchar *ip_address = (gchar *) user_data;

	int limit = abs(GPOINTER_TO_INT(g_hash_table_lookup(rate_limits, ip_address)));
	g_hash_table_replace(rate_limits, ip_address, GINT_TO_POINTER(-limit));

	/* We can't free the ip address yet, since it still needs it for the
	 * key. TODO: How not to leak? */

	return FALSE;
}

static void
sapphire_got_internal(gchar *frame)
{
	JsonNode *root = json_parse_to_root(frame);

	if (root == NULL) {
		printf("NULL root, ignoring\n");
		return;
	}

	JsonObject *obj = json_node_get_object(root); 

	const gchar *op = json_object_get_string_member(obj, "op");

	if (g_strcmp0(op, "icon") == 0) {
		SapphireIcon *icon = g_new0(SapphireIcon, 1);

		const gchar *name = json_object_get_string_member(obj, "name");
		const gchar *base64 = json_object_get_string_member(obj, "base64");

		icon->extension = json_object_get_string_member(obj, "ext");
		icon->data = g_base64_decode(base64, &icon->size);

		g_hash_table_insert(username_to_icon, g_strdup(name), icon);
	} else {
		printf("Unknown op %s\n", op);
		return;
	}
}

static void sapphire_got_line(GObject *source_object, GAsyncResult *res, gpointer user_data);

static void
sapphire_read_line(Connection *conn)
{
	g_data_input_stream_read_line_async(conn->distream, G_PRIORITY_DEFAULT, NULL, sapphire_got_line, conn);
}

static void
sapphire_got_line(GObject *source_object,
                        GAsyncResult *res,
                        gpointer user_data)
{
	Connection *conn = (Connection *) user_data;

	GError *err = NULL;
	gsize len;
	char *data = g_data_input_stream_read_line_finish_utf8(G_DATA_INPUT_STREAM(source_object), res, &len, &err);

	if (err || !data) {
		/* Borp, error -- disconnect */
		printf("Disconnecting %s...\n", err ? err->message : "");

		if (data)
			g_free(data);

		sapphire_close_connection(conn);

		return;
	}

	/* Check the first character. If it's >, this is an internal packet. Otherwise, pass it along */

	if (data[0] == '>') {
		sapphire_got_internal(data + 1);
	} else {
		sapphire_send_raw_packet(conn, data);
	}

	g_free(data);

	sapphire_read_line(conn);
}

static gboolean
sapphire_try_login(Connection *conn, const char *username, const char *attempted_hash);

static void
sapphire_send_rate_limit(Connection *conn, int ms);

static void
soup_ws_data(SoupWebsocketConnection *self,
               gint                     type,
               GBytes                  *message,
               gpointer                 user_data)
{
	struct Connection *conn = (struct Connection *) user_data;

	const gchar *frame = (const gchar *) g_bytes_get_data(message, NULL);

	/* The message should be interpreted as JSON, decode that here */

	JsonParser *parser = json_parser_new();

	if (!json_parser_load_from_data(parser, frame, -1, NULL)) {
		fprintf(stderr, "Error parsing response: ...\n");
		return;
	}

	JsonNode *root = json_parser_get_root(parser);

	if (root == NULL) {
		printf("NULL root, ignoring\n");
		return;
	}

	/* How to proceed depends if we're authenticated or not. If we are,
	 * this is a standard client message, ready to be parsed, relayed, and
	 * actuated. If we are not, this is an authentication message (by
	 * definition -- otherwise they get booted to penalize credential
	 * attacks */

	if (conn->is_authenticated) {
		/* Forward the authenticated packet */

		if (!conn->proxy_connection) {
			printf("No proxy\n");
			return;
		}

		GError *gerror;
		GOutputStream *ostream = g_io_stream_get_output_stream (G_IO_STREAM (conn->proxy_connection));

		g_output_stream_write_all(ostream, frame, strlen(frame), NULL, NULL, &gerror);
		char end = '\n';
		g_output_stream_write(ostream, &end, 1, NULL, &gerror);

	} else {
		JsonObject *obj = json_node_get_object(root);

		gboolean success = FALSE;
		const char *username;

		if (obj) {
			username = json_object_get_string_member(obj, "username");

			const char *passwordHash = json_object_get_string_member(obj, "passwordHash");

			if (username && passwordHash) {
				success = sapphire_try_login(conn, username, passwordHash);
			}
		}

		if (!success) {
			/* Slow down future attempts for rate limiting */
			int limit = abs(GPOINTER_TO_INT(g_hash_table_lookup(rate_limits, conn->ip_address)));

			if (limit == 0) {
				/* Initial rate limit of 320ms, grows exponentially by two's */
				limit = 1;
			} 

			/* Exponential rate limit growth for repeat offenders */
			limit *= 2;

			/* Store that limit */
			g_hash_table_insert(rate_limits, g_strdup(conn->ip_address), GINT_TO_POINTER(limit));

			/* Create a timeout to restore their access */
			int milliseconds = limit * 160;
			g_timeout_add(milliseconds, sapphire_restore_rate_limit, g_strdup(conn->ip_address));

			/* Tell the client how long we're rate limiting them for */
			sapphire_send_rate_limit(conn, milliseconds);

			/* Eject */
			const char *error = "Authentication error";
			soup_websocket_connection_close(conn->connection, SOUP_WEBSOCKET_CLOSE_POLICY_VIOLATION, error);

			return;
		} else {
			/* Successful login - so get rid of the rate limit */

			g_hash_table_replace(rate_limits, g_strdup(conn->ip_address), GINT_TO_POINTER(0));
		}



		/* To authenticate, set the flag and add us to the list */
		conn->is_authenticated = TRUE;	

		/* Remove size limit -- needed for avatar upload, etc. TODO: Is this risky? */
		soup_websocket_connection_set_max_incoming_payload_size(conn->connection, 0);

		authenticated_connections = g_list_append(authenticated_connections, conn);

		/* ...and connect to the appropriate backend's socket */

		gchar *socket_path = multi_user ? g_strdup_printf("./accounts/%s/sockpuppet", username) : g_strdup("./sockpuppet");

		GSocketClient *client = g_socket_client_new();
		GSocketAddress *addr = g_unix_socket_address_new(socket_path);
		GSocketConnection *proxy_conn = g_socket_client_connect(client, G_SOCKET_CONNECTABLE(addr), NULL, NULL);

		if (!proxy_conn) {
			fprintf(stderr, "Failed to proxy\n");

			JsonObject *data = json_object_new();
			json_object_set_string_member(data, "op", "proxyerror");
			sapphire_send(conn, data);
			json_object_unref(data);
			return;
		}
		
		conn->proxy_connection = g_object_ref(proxy_conn);

		GInputStream *istream = g_io_stream_get_input_stream (G_IO_STREAM (conn->proxy_connection));
		conn->distream = g_data_input_stream_new(istream);

		/* Start the async */
		sapphire_read_line(conn);
	}
}

static void
soup_ws_error(SoupWebsocketConnection *self,
		GError *gerror,
               gpointer                 user_data)
{
	fprintf(stderr, "WS Error\n");
}

static void
sapphire_close_connection(Connection *conn)
{
	/* Free connection */
	if (conn->ip_address) {
		g_free(conn->ip_address);
		conn->ip_address = NULL;
	}

	/* Disconnect the proxy-half of the connection */
	if (conn->proxy_connection) {
		g_io_stream_close(G_IO_STREAM(conn->proxy_connection), NULL, NULL);
		conn->proxy_connection = NULL;
	}

	/* Splice the socket out of the authenticated list, so we no longer
	 * attempt to broadcast to it */

	authenticated_connections = g_list_remove(authenticated_connections, conn);
}

static void
soup_ws_closed(SoupWebsocketConnection *self,
               gpointer                 user_data)
{
	Connection *conn = (Connection *) user_data;
	sapphire_close_connection(conn);
}

static void
soup_ws_callback(SoupServer *server,
			SoupWebsocketConnection *connection,
			const char *path,
			SoupClientContext *client,
			gpointer user_data)
{
	/* Figure out who we're talking to */
	GSocketAddress *socket_address = soup_client_context_get_remote_address(client);

	GSocketFamily family = g_socket_address_get_family(socket_address);

	if ((family != G_SOCKET_FAMILY_IPV4) && (family != G_SOCKET_FAMILY_IPV6)) {
		/* Should be unreachable */
		fprintf(stderr, "Non-IP socket?\n");
		return;
	}

	GInetAddress *inet_address = g_inet_socket_address_get_address((GInetSocketAddress *) socket_address);
	gchar *addr = g_inet_address_to_string(inet_address);

	/* Allocate a connection object for us and fill it in */
	Connection *conn = g_new0(Connection, 1);

	conn->is_authenticated = FALSE;

	/* Save the connection.
	 * IMPORTANT: Reference counting is necessary to keep the connection
	 * alive. No idea why this isn't documented anywhere, xxx
	 */

	conn->connection = g_object_ref(connection);

	/* Save the IP for ratelimiting */
	conn->ip_address = addr;

	/* Check for rate limiting, for that matter */
	int rate_limit = GPOINTER_TO_INT(g_hash_table_lookup(rate_limits, conn->ip_address));

	if (rate_limit > 0) {
		/* Violation -- eject */
		const char *error = "Rate limit violation";
		soup_websocket_connection_close(conn->connection, SOUP_WEBSOCKET_CLOSE_POLICY_VIOLATION, error);
		printf("Ejecting rate limit violation from %s\n", conn->ip_address);
		return;
	}

	/* Subscribe to the various signals */

	g_signal_connect(connection, "message", G_CALLBACK(soup_ws_data), conn);
	g_signal_connect(connection, "closed", G_CALLBACK(soup_ws_closed), conn);
	g_signal_connect(connection, "error", G_CALLBACK(soup_ws_error), conn);
}

static void
sapphire_init_websocket(gboolean secure)
{
	GError *error = NULL;

	SoupServer *soup = soup_server_new(NULL, NULL);

	if (secure) {
		gchar *key_file = "key.pem";
		gchar *cert_file = "cert.pem";

		if (!soup_server_set_ssl_cert_file(soup, cert_file, key_file, &error)) {
			printf("Error setting SSL certificate\n");
			printf("Msg: %s\n", error->message);
			exit(1);
		}
	} else {
		/* Be a little scary */
		fprintf(stderr, "* * *\n\nWARNING: RUNNING IN PLAIN HTTP MODE!!! DO NOT DEPLOY IN PRODUCTION!!!\n\n* * *\n\n");
	}


	/* Initialize icon endpoint */
	soup_server_add_handler(soup, "/icon/", soup_icon_callback, NULL, NULL);

	/* WebSocket entrypoint */
	char *protocols[] = { "binary", NULL };
	soup_server_add_websocket_handler(soup, "/ws", NULL, protocols, soup_ws_callback, NULL, NULL);

	if (!soup_server_listen_all(soup, WS_PORT, secure ? SOUP_SERVER_LISTEN_HTTPS : 0, &error)) {
		fprintf(stderr, "Error listening in soup\n");
		fprintf(stderr, "Msg: %s\n", error->message);
		exit(1);
	}
}



#include "secure-compare-64.h"

static gboolean
sapphire_try_login(Connection *conn, const char *username, const char *attempted_hash)
{
	gboolean success = FALSE;

	/* Find the appropriate account */

	SapphireAccount *account = g_hash_table_lookup(username_to_account, username);

	if (!account) {
		printf("Bad account\n");
		return FALSE;
	}

	/* Valid hashes must be 64 bytes (32 bytes of binary -> 64 of hex) */

	if (strlen(attempted_hash) == 64 && strlen(account->password_hash) == 64) {
		/* Fetch the SHA-256 secret */
		success = secure_compare_64(attempted_hash, account->password_hash);
	}

	/* TODO: Rate limit on failure */

	/* Packet indicating status */

	if (success) {
		JsonObject *data = json_object_new();
		json_object_set_string_member(data, "op", "authsuccess");
		sapphire_send(conn, data);
		json_object_unref(data);
	}

	return success;
}

static void
sapphire_send_rate_limit(Connection *conn, int ms)
{
	JsonObject *data = json_object_new();
	json_object_set_string_member(data, "op", "ratelimit");
	json_object_set_int_member(data, "milliseconds", ms);
	sapphire_send(conn, data);
	json_object_unref(data);
}

/* A buddy joined in a room we're subscribed to -- but that doesn't mean the
 * client needs to know. Only send the joined event to clients that have opened
 * the corresponding conversation */

GList *authenticated_connections;

int main(int argc, char *argv[])
{
	GMainLoop *loop;

#ifdef _WIN32
	g_thread_init(NULL);
#endif

	g_set_prgname("Sapphire-Proxy");
	g_set_application_name("Sapphire-Proxy");

	loop = g_main_loop_new(NULL, FALSE);
	g_main_loop_ref(loop);

	/* Initialize icon database */
	username_to_icon = g_hash_table_new(g_str_hash, g_str_equal);
	rate_limits = g_hash_table_new(g_str_hash, g_str_equal);
	username_to_account = g_hash_table_new(g_str_hash, g_str_equal);

	gboolean secure = (argc >= 2) && (g_strcmp0(argv[1], "--production") == 0);

	/* Only allow multi-user mode if we're TLS-encrypted */
	multi_user = secure;

	const gchar *database_path = (argc >= 3) ? argv[2] : "./sapphire-accounts.json";
	sapphire_load_accounts(database_path);

	sapphire_init_websocket(secure);

	g_main_context_iteration(g_main_loop_get_context(loop), FALSE);
	
	g_main_loop_run(loop);

	return 0;
}
